Import reduce naive from burn #314
Conversation
| let mut offset_input = 0; | ||
| for axis in 0..input.rank() { | ||
| let coordinate = (ABSOLUTE_POS / output.stride(axis)) % output.shape(axis); | ||
| offset_input += coordinate * input.stride(axis); |
There was a problem hiding this comment.
How come the if axis != dim was removed? Was it useless?
There was a problem hiding this comment.
The shape of the dim axis in output is always 1, so for that particular axis, the coordinate is always 0. Thus, it has a null effect on offset_input.
There was a problem hiding this comment.
SHould I add a comment?
There was a problem hiding this comment.
No, I was just making sure nothing was lost from the original
crates/cubecl-reduce/src/naive.rs
Outdated
| } | ||
|
|
||
| fn accumulate(accumulator: &mut Self::Accumulator, current_value: Line<EI>, i: u32) { | ||
| // FIX: There is an issue when line_size is 1 on wgpu. |
There was a problem hiding this comment.
You could do if comptime!(current_value.size() == 1) and have a different logic without indexes
| .collect() | ||
| } | ||
|
|
||
| fn random_input_values<F: Float>(&self) -> Vec<F> { |
There was a problem hiding this comment.
Quite a duplication from matmul test_utils. This would be a good case for cubecl-std, under a test cfg. Also, the matmul one has a seed as argument which can prevent bug proneness when we have several inputs. I think we should have an Option and use 123456789 when None. Anyway, not a blocker for this PR.
There was a problem hiding this comment.
I can create an issue?
There was a problem hiding this comment.
Yeah, maybe an issue about creating cubecl-std with common stuff between reduce and matmul.
This is the first PR to import the reduction algorithms implemented by @wingertge in Burn.
Currently, this only import the naive implementation. The others will follow soon. I also added a bunch of tests, support for lines and refactor the code a little bit. However, the logic of the implementation follow the one from Burn.
Notice that the tests currently ignore ArgMax and ArgMin since I have some issues with Line that most be solved first.